from types import SimpleNamespace

from wenet.zipformer.init_encoder_embedding import init_encoder_embedding
from wenet.zipformer.init_decoder_attention import init_decoder4zipformer
from wenet.zipformer.init_zipformer import init_zipformer2_encoder
from wenet.zipformer.asr_model import K2AsrModel
from wenet.zipformer.config import CONFIG_ZIPFORMER



def init_k2model():
    config_model = CONFIG_ZIPFORMER
    params = SimpleNamespace(**config_model)
    encoder_embed = init_encoder_embedding()
    encoder = init_zipformer2_encoder(config_model)
    attention_decoder = init_decoder4zipformer(config_model)
    model = K2AsrModel(
        encoder_embed=encoder_embed,
        encoder=encoder,
        decoder=None,
        joiner=None,
        attention_decoder=attention_decoder,
        encoder_dim=max(params.encoder_dim),
        decoder_dim=params.decoder_dim,
        vocab_size=params.vocab_size,
        use_transducer=False,
        use_ctc=True,
        use_attention_decoder=True,
    )
    return model